import numpy as np


class ANFIS:
    def __init__(self, n_inputs, n_rules):
        self.n_inputs = n_inputs  # 输入变量数量
        self.n_rules = n_rules  # 模糊规则数量
        # 初始化前提（模糊）参数，均值和标准差
        self.premise_params = np.random.rand(n_inputs, n_rules, 2)
        # 初始化结论参数，y = p * x1 + q * x2 + r
        self.consequence_params = np.random.rand(n_rules, n_inputs + 1)

    def gaussmf(self, x, mean, sigma):
        """高斯隶属函数"""
        return np.exp(-((x - mean) ** 2) / (2 * sigma ** 2))

    def forward(self, x):
        """前向传播，计算每个规则的激活度并输出预测值"""
        # 1. 计算隶属度
        membership_values = np.zeros((self.n_rules, self.n_inputs))
        for i in range(self.n_inputs):
            for j in range(self.n_rules):
                mean, sigma = self.premise_params[i, j]
                membership_values[j, i] = self.gaussmf(x[i], mean, sigma)

        # 2. 计算规则激活强度
        firing_strengths = np.prod(membership_values, axis=1)
        firing_strengths /= np.sum(firing_strengths)  # 归一化

        # 3. 计算规则的线性输出
        rule_outputs = np.dot(self.consequence_params[:, :-1], x) + self.consequence_params[:, -1]

        # 4. 计算最终输出
        output = np.dot(firing_strengths, rule_outputs)
        return output, firing_strengths, rule_outputs

    def train(self, X, y, epochs=100, lr=0.01):
        """训练ANFIS模型，使用梯度下降调整参数"""
        for epoch in range(epochs):
            total_loss = 0
            for i in range(X.shape[0]):
                x_sample, y_true = X[i], y[i]

                # 前向传播
                y_pred, firing_strengths, rule_outputs = self.forward(x_sample)

                # 计算损失
                loss = (y_pred - y_true) ** 2
                total_loss += loss

                # 反向传播 - 更新参数
                for j in range(self.n_rules):
                    # 更新结论参数
                    error = y_pred - y_true
                    self.consequence_params[j, :-1] -= lr * error * firing_strengths[j] * x_sample
                    self.consequence_params[j, -1] -= lr * error * firing_strengths[j]

                    # 更新前提参数（隶属度参数）
                    for i in range(self.n_inputs):
                        mean, sigma = self.premise_params[i, j]
                        x = x_sample[i]
                        grad_mean = firing_strengths[j] * (x - mean) / (sigma ** 2) * self.gaussmf(x, mean, sigma)
                        grad_sigma = firing_strengths[j] * ((x - mean) ** 2) / (sigma ** 3) * self.gaussmf(x, mean,
                                                                                                           sigma)
                        self.premise_params[i, j, 0] -= lr * error * grad_mean
                        self.premise_params[i, j, 1] -= lr * error * grad_sigma

            # 打印每个epoch的损失
            print(f"Epoch {epoch + 1}/{epochs}, Loss: {total_loss / X.shape[0]}")

    def predict(self, X):
        """预测函数"""
        y_pred = []
        for x in X:
            output, _, _ = self.forward(x)
            y_pred.append(output)
        return np.array(y_pred)


# 示例用法
# 创建训练数据
X = np.array([[1.0, 2.0], [2.0, 1.0], [2.5, 2.5], [3.0, 3.0]])
y = np.array([3.0, 3.5, 4.0, 5.0])

# 实例化ANFIS模型
anfis = ANFIS(n_inputs=2, n_rules=2)

# 训练模型
anfis.train(X, y, epochs=50, lr=0.01)

# 预测
predictions = anfis.predict(X)
print("Predictions:", predictions)
